import os
import argparse

"""
extract images from yuv videos
"""

parser = argparse.ArgumentParser()
parser.add_argument("-f", "--seq_folder", type = str)
parser.add_argument("-sw", "--seq_width", type = str)
parser.add_argument("-sh", "--seq_height", type = str)

parser.add_argument("-o", "--save_folder", type = str, default = "./uncompressed")
parser.add_argument("--format", type = str, default = "png")

args = parser.parse_args()

seq_folder = args.seq_folder
sw, sh = args.seq_width, args.seq_height

save_folder = args.save_folder
format = args.format


if not os.path.exists(save_folder):
    os.makedirs(save_folder)


for seq in os.listdir(seq_folder):
    if not seq.endswith(".yuv"):
        continue

    input = os.path.join(seq_folder, seq)
    output_folder = os.path.join(save_folder, seq[:-4])
    
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    options = [
        "ffmpeg",
        "-pix_fmt", "yuv420p",
        "-s:v", f"{sw}x{sh}",
        "-i", input,
        "-f", "image2",
        f"{output_folder}/im%04d.{format}"
    ]

    cmd = ""
    for opt in options:
        cmd += (opt + " ").replace("\\", "/")

    print(cmd)
    os.system(cmd)

